# -*- coding: utf-8 -*-  
'''
初始化字典，并写入磁盘
顺便检测max_len是多少

@author: luoyi
Created on 2021年3月18日
'''
import os
import pickle

import data.dataset as ds_couplet
import utils.conf as conf
import utils.dictionaries_bert as dictionaries


#    跑一遍全体数据集，初始化字典
train_reader = ds_couplet.read_couplet_iterator(in_path=conf.DATASET.get_in_train(), 
                                                out_path=conf.DATASET.get_label_train(), 
                                                count=-1)

val_reader = ds_couplet.read_couplet_iterator(in_path=conf.DATASET.get_in_val(), 
                                              out_path=conf.DATASET.get_label_val(), 
                                              count=-1)

test_reader = ds_couplet.read_couplet_iterator(in_path=conf.DATASET.get_in_test(), 
                                               out_path=conf.DATASET.get_label_test(), 
                                               count=-1)

max_len = 0

print('初始化字典，遍历训练集数据...')
for first_couplet, second_couplet in train_reader:
    max_len = max(len(first_couplet), len(second_couplet), max_len)
    dictionaries.word2idx_slist(first_couplet)
    dictionaries.word2idx_slist(second_couplet)
    pass
print('初始化字典，训练数据集遍历完毕，字典长度:', dictionaries.dict_size(), ' max_len:', max_len)
print('初始化字典，遍历验证集数据：')
for first_couplet, second_couplet in train_reader:
    max_len = max(len(first_couplet), len(second_couplet), max_len)
    dictionaries.word2idx_slist(first_couplet)
    dictionaries.word2idx_slist(second_couplet)
    pass
print('初始化字典，验证数据集遍历完毕，字典长度:', dictionaries.dict_size(), ' max_len:', max_len)
print('初始化字典，遍历测试集数据：')
for first_couplet, second_couplet in train_reader:
    max_len = max(len(first_couplet), len(second_couplet), max_len)
    dictionaries.word2idx_slist(first_couplet)
    dictionaries.word2idx_slist(second_couplet)
    pass
print('初始化字典，测试数据集遍历完毕，字典长度:', dictionaries.dict_size(), ' max_len:', max_len)


#    写入磁盘
fdir = conf.TEXT.get_dictionaries_path()
word2dict_path = fdir + '/word2dict_bert.pkl'
dict2word_path = fdir + '/dict2word_bert.pkl'
#    检测目录是否存在
conf.mkdir_ifnot_exises(fdir)
#    存在同名文件先删除
if (os.path.exists(word2dict_path)): os.remove(word2dict_path)
if (os.path.exists(dict2word_path)): os.remove(dict2word_path)
#    写
with open(word2dict_path, mode='wb') as file_word2dict:
    pickle.dump(dictionaries.word_to_index, file_word2dict)
with open(dict2word_path, mode='wb') as file_dict2word:
    pickle.dump(dictionaries.index_to_word, file_dict2word)
    
print('max_len:', max_len)